package main

import (
    "github.com/tecbot/gorocksdb"
    "path"
    "fmt"
    "sync/atomic"
    "time"
)

var BasePath = "./"

type RocksdbStore struct {
    tableDef      TableDefine
    partition     Partition
    lastReadTime  *int64
    lastWriteTime *int64
    lastOpenTime  *int64
    db            *gorocksdb.DB
    wo            *gorocksdb.WriteOptions
    ro            *gorocksdb.ReadOptions
    opt           *gorocksdb.Options
}


func NewRocksdbStore(tableDef TableDefine, partition Partition) (*RocksdbStore, error) {
    if tableDef.DtLength != len(partition.dt) {
        return nil, fmt.Errorf("dt length not match")
    }
    opt := gorocksdb.NewDefaultOptions()
    rateLimiter := gorocksdb.NewRateLimiter(1024, 100*1000, 10)
    opt.SetRateLimiter(rateLimiter)
    opt.SetCreateIfMissing(true)
    ro := gorocksdb.NewDefaultReadOptions()
    wo := gorocksdb.NewDefaultWriteOptions()
    return &RocksdbStore{
        wo:        wo,
        ro:        ro,
        opt:       opt,
        tableDef:  tableDef,
        partition: partition,
    }, nil
}

func (tbs *RocksdbStore) GetTableDef() TableDefine {
    return tbs.tableDef
}
func (tbs *RocksdbStore) GetPartition() Partition {
    return tbs.partition
}

func (tbs *RocksdbStore) Write(rows []Row) error {
    wb := gorocksdb.NewWriteBatch()
    defer wb.Destroy()
    for _, row := range rows {
        wb.Put([]byte(row.Key), []byte(row.Value))
        indexes, err := tbs.tableDef.BuildIndex(row)
        if err != nil {
            return err
        }
        for _, idx := range indexes {
            wb.Put([]byte(idx), []byte{})
        }
    }
    atomic.StoreInt64(tbs.lastWriteTime,time.Now().Unix())
    return tbs.db.Write(tbs.wo, wb)
}

func (tbs *RocksdbStore) LastActiveTime() (int64, int64, int64) {
    return atomic.LoadInt64(tbs.lastOpenTime), atomic.LoadInt64(tbs.lastReadTime), atomic.LoadInt64(tbs.lastWriteTime)
}

func (tbs *RocksdbStore) Query(k string) (string, error) {
    v, err := tbs.db.Get(tbs.ro, []byte(k))
    defer v.Free()
    if err != nil {
        return "", err
    }
    atomic.StoreInt64(tbs.lastReadTime,time.Now().Unix())
    return string(v.Data()), nil
}

func (tbs *RocksdbStore) QueryKeys(condition QueryCondition) ([]string, error) {
    var keys [][]string
    if condition.keyPrefix != "" {
        keys = append(keys, tbs.scan([]byte(condition.keyPrefix), false))
    }
    pres, err := tbs.tableDef.BuildQueryIndexPrefixes(condition)
    if err != nil {
        return nil, err
    }
    for _, p := range pres {
        k := tbs.scan([]byte(p), true)
        keys = append(keys, k)
    }
    atomic.StoreInt64(tbs.lastReadTime,time.Now().Unix())
    return Keys(keys), nil
}

func (tbs *RocksdbStore) QueryValues(condition QueryCondition) (map[string]string, error) {
    ret := make(map[string]string)
    keys, err := tbs.QueryKeys(condition)
    if err != nil {
        return nil, err
    }
    for _, key := range keys {
        value, err := tbs.Query(key)
        if err != nil {
            return nil, err
        }
        ret[key] = value
    }
    atomic.StoreInt64(tbs.lastReadTime,time.Now().Unix())
    return ret, nil
}

func (tbs *RocksdbStore) scan(prefix []byte, isIndex bool) (ret []string) {
    iter := tbs.db.NewIterator(tbs.ro)
    defer iter.Close()
    for iter.Seek(prefix); iter.ValidForPrefix(prefix); iter.Next() {
        raw := iter.Key().Data()
        if isIndex {
            raw = tbs.tableDef.ExtractKey(raw)
        }
        if len(raw) > 0 {
            ret = append(ret, string(raw))
        }
    }
    return nil
}

func (tbs *RocksdbStore) Open() error {
    dir := path.Join(BasePath, tbs.tableDef.Path(), tbs.partition.StoreName())
    db, err := gorocksdb.OpenDb(tbs.opt, dir)
    if err != nil {
        return err
    }
    tbs.db = db
    atomic.StoreInt64(tbs.lastOpenTime,time.Now().Unix())
    return nil
}

func (tbs *RocksdbStore) Close() error {
    if tbs.db != nil {
        tbs.db.Close()
        tbs.db = nil
    }
    return nil
}

func (tbs *RocksdbStore) IsActive() bool {
    return tbs.db == nil
}
